# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： bbox_transform
Description :
Author : 'li'
date： 2022/6/16
Change Activity:
2022/6/16:
-------------------------------------------------
"""
import torch

from ml.cv.util.bbox_transform import box_xyxy2cxcywh


class BoxXyxyToCxcywh:
    def __init__(self, input_key='boxes', img_key='img', output_key='boxes'):
        """

        Args:
            input_key:
            output_key:
        """
        self.input_key = input_key
        self.output_key = output_key
        self.img_key = img_key

    def __call__(self, data):
        boxes = data[self.input_key]
        image = data[self.img_key]
        h, w = image.shape[-2:]
        boxes = box_xyxy2cxcywh(boxes)
        boxes = boxes / torch.tensor([w, h, w, h], dtype=torch.float32)
        data[self.output_key] = boxes
        return data
